package com.hanliukui.ai.langchain4j.infrastructure;

import com.hanliukui.ai.langchain4j.infrastructure.domain.repository.ChatMessagesPoRepository;
import dev.langchain4j.community.model.dashscope.QwenChatModel;
import dev.langchain4j.community.model.dashscope.QwenEmbeddingModel;
import dev.langchain4j.community.model.dashscope.QwenStreamingChatModel;
import dev.langchain4j.community.store.embedding.redis.RedisEmbeddingStore;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.chat.StreamingChatModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

@Configuration
public class ModelConfig {

    public static final String DASH_SCOPE_API_KEY = System.getenv("AI_DASHSCOPE_API_KEY");

    @Bean
    public ChatModel chatModel() {
        return new QwenChatModel.QwenChatModelBuilder()
                .apiKey(DASH_SCOPE_API_KEY)
                .modelName("qwen-plus")
                .temperature(0.7f)
                .maxTokens(4096)
                .build();
    }

    @Bean
    public StreamingChatModel streamingChatModel() {
        return new QwenStreamingChatModel.QwenStreamingChatModelBuilder()
                .apiKey(DASH_SCOPE_API_KEY)
                .modelName("qwen-plus")
                .temperature(0.7f)
                .maxTokens(4096)
                .build();
    }

    @Bean
    public EmbeddingModel embeddingModel() {
        return new QwenEmbeddingModel
                .QwenEmbeddingModelBuilder()
                .apiKey(DASH_SCOPE_API_KEY)
                .modelName("text-embedding-v3")
                .dimension(1024)
                .build();
    }

    @Bean
    public ChatMemoryStore chatMemoryStore(ChatMessagesPoRepository repository) {
        return new JpaChatMemoryStore(repository);
    }

    @Bean
    public EmbeddingStore<TextSegment> embeddingStore() {
//        return new InMemoryEmbeddingStore<>();
        // 使用redis做向量存储
        return new RedisEmbeddingStore.Builder()
                .host("192.168.18.218")
                .port(6379)
                .user("default")
                .password("123456")
                .indexName("langchain_embedding_store")
                .prefix("langchain_embedding_store:")
                .dimension(1024)
                .build();
    }


}
